package com.chenl.provider.DEMOS.currLimit限流;

import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.TimeUnit;

@Slf4j
@Aspect
@Configuration
public class LimitAspect {
    @Resource(name = "MyRedisTemplate")
    private RedisTemplate redisTemplate;
    private static final DefaultRedisScript<Long> RATE_LIMIT;

    //传入限流key KEYS[1],限流个数 ARGV[1],过期时间(秒) ARGV[2]
    static {
        StringBuilder lua = new StringBuilder();
        lua.append("local key = KEYS[1]\n"                             //1. key - 令牌桶的 key
                + "local intervalPerTokens = tonumber(ARGV[1])\n"     //2. intervalPerTokens - 生成令牌的间隔(ms)
                + "local curTime = tonumber(ARGV[2])\n"               //3. curTime - 当前时间
                + "local initTokens = tonumber(ARGV[3])\n"            //4. initTokens - 令牌桶初始化的令牌数
                + "local bucketMaxTokens = tonumber(ARGV[4])\n"       //5. bucketMaxTokens - 令牌桶的上限
                + "local resetBucketInterval = tonumber(ARGV[5])\n"   //6. resetBucketInterval - 重置桶内令牌的时间间隔
                + "local bucket = redis.call('hgetall', key)\n"       //7. currentTokens - 当前桶内令牌数
                + "local currentTokens\n"                             //8. bucket - 当前 key 的令牌桶对象
                //若当前桶未初始化,先初始化令牌桶
                + "if table.maxn(bucket) == 0 then\n" +
                //初始桶内令牌
                "currentTokens = initTokens\n" +
                //置桶最近的填充时间是当前
                "redis.call('hset', key, 'lastRefillTime', curTime)\n" +
                //初始化令牌桶的过期时间, 设置为间隔的
                "redis.call('expire', key, resetBucketInterval * 1.5)\n" +
                //若桶已初始化,开始计算桶内令牌
                //为什么等于 4 ? 因为有两对 field, 加起来长度是 4\n" +
                //{ \"lastRefillTime(上一次更新时间)\",\"curTime(更新时间值)\",\"tokensRemaining(当前保留的令牌)\",\"令牌数\" }\n
                "elseif table.maxn(bucket) == 4 then\n" +
                //上次填充时间
                "local lastRefillTime = tonumber(bucket[2])\n" +
                //剩余的令牌数
                "local tokensRemaining = tonumber(bucket[4])\n" +
                //当前时间大于上次填充时间
                "if curTime > lastRefillTime then\n" +
                //拿到当前时间与上次填充时间的时间间隔
                //举例理解: curTime = 2620 , lastRefillTime = 2000, intervalSinceLast = 620
                "local intervalSinceLast = curTime - lastRefillTime\n" +
                //如果当前时间间隔 大于 令牌的生成间隔
                //举例理解: intervalSinceLast = 620, resetBucketInterval = 1000
                "if intervalSinceLast > resetBucketInterval then\n" +
                //将当前令牌填充满
                "currentTokens = initTokens\n" +
                //更新重新填充时间
                "redis.call('hset', key, 'lastRefillTime', curTime)\n" +
                //如果当前时间间隔 小于 令牌的生成间隔
                "else\n" +
                //可授予的令牌 = 向下取整数( 上次填充时间与当前时间的时间间隔 / 两个令牌许可之间的时间间隔 )
                //举例理解 : intervalPerTokens = 200 ms , 令牌间隔时间为 200ms
                //intervalSinceLast = 620 ms , 当前距离上一个填充时间差为 620ms
                //grantedTokens = 620/200 = 3.1 = 3
                "local grantedTokens = math.floor(intervalSinceLast / intervalPerTokens)\n" +
                //可授予的令牌 > 0 时
                //举例理解 : grantedTokens = 620/200 = 3.1 = 3
                "if grantedTokens > 0 then\n" +
                //生成的令牌 = 上次填充时间与当前时间的时间间隔 % 两个令牌许可之间的时间间隔
                //举例理解 : padMillis = 620%200 = 20
                //          curTime = 2620
                //          curTime - padMillis = 2600
                "local padMillis = math.fmod(intervalSinceLast, intervalPerTokens)\n" +
                //将当前令牌桶更新到上一次生成时间
                "redis.call('hset', key, 'lastRefillTime', curTime - padMillis)\n" +
                "end\n" +
                //更新当前令牌桶中的令牌数
                //Math.min(根据时间生成的令牌数 + 剩下的令牌数, 桶的限制) => 超出桶最大令牌的就丢弃
                "currentTokens = math.min(grantedTokens + tokensRemaining, bucketMaxTokens)\n" +
                "end\n" +
                "else\n" +
                //如果当前时间小于或等于上次更新的时间, 说明刚刚初始化, 当前令牌数量等于桶内令牌数
                //不需要重新填充
                "currentTokens = tokensRemaining\n" +
                "end\n" +
                "end\n" +
                //如果当前桶内令牌小于 0,抛出异常
                "assert(currentTokens >= 0)\n" +
                //如果当前令牌 == 0 ,更新桶内令牌, 返回 0
                "if currentTokens == 0  then\n" +
                "redis.call('hset', key, 'tokensRemaining', currentTokens)\n" +
                "return 0\n" +
                "else\n" +
                //如果当前令牌 大于 0, 更新当前桶内的令牌 -1 , 再返回当前桶内令牌数
                "redis.call('hset', key, 'tokensRemaining', currentTokens - 1)\n" +
                "return currentTokens\n" +
                "end\n");
        RATE_LIMIT = new DefaultRedisScript<>(lua.toString(), Long.class);
    }

    @Around("@annotation(com.chenl.provider.DEMOS.currLimit限流.RateLimit)")
    public Object interceptor(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);
        if (rateLimit != null) {
            HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
            String ipAddress = getIpAddr(request);
            StringBuffer stringBuffer = new StringBuffer();
            if (rateLimit.type().equals(AnnotationType.RATE_LIMIT_IP)) {//根据ip限流
                stringBuffer.append(ipAddress).append("_")
                        .append(targetClass.getName()).append("_")
                        .append(method.getName()).append("_")
                        .append(rateLimit.key());
            } else if (rateLimit.type().equals(AnnotationType.RATE_LIMIT_INTERFACE)) {//根据接口限流
                stringBuffer
                        .append(targetClass.getName()).append("_")
                        .append(method.getName()).append("_")
                        .append(rateLimit.key());
            }
            List<String> keys = new ArrayList<>();
            keys.add(stringBuffer.toString().replace(".", "").replace(" ", ""));
            long curTime = System.currentTimeMillis();
            long intervalPerTokens = rateLimit.intervalPerTokens();
            log.info("LimitAspect---KEYS:{},ARGV[1]:{},ARGV[2]:{},ARGV[3]:{},ARGV[4]:{},ARGV[5]:{}", JSONObject.toJSONString(keys)
                    , intervalPerTokens
                    , curTime
                    , rateLimit.bucketMaxTokens()
                    , rateLimit.bucketMaxTokens()
                    , rateLimit.resetBucketInterval());
            if (rateLimit.intervalPerTokens() == 0) {//没设置则自行进行计算令牌投放间隔
                int time = rateLimit.time();
                int count = rateLimit.count();
                TimeUnit timeUnit = rateLimit.timeType();
                BigDecimal ms = BigDecimal.valueOf(1000);
                if (timeUnit == TimeUnit.SECONDS) {
                    intervalPerTokens = Long.parseLong(BigDecimal.valueOf(time).multiply(ms).divide(BigDecimal.valueOf(count)).toString());
                } else if (timeUnit == TimeUnit.MINUTES) {
                    ms = ms.multiply(BigDecimal.valueOf(60));
                    intervalPerTokens = Long.parseLong(BigDecimal.valueOf(time).multiply(ms).divide(BigDecimal.valueOf(count)).toString());
                } else if (timeUnit == TimeUnit.HOURS) {
                    ms = ms.multiply(BigDecimal.valueOf(3600));
                    intervalPerTokens = Long.parseLong(BigDecimal.valueOf(time).multiply(ms).divide(BigDecimal.valueOf(count)).toString());
                } else if (timeUnit == TimeUnit.DAYS) {
                    ms = ms.multiply(BigDecimal.valueOf(86400));
                    intervalPerTokens = Long.parseLong(BigDecimal.valueOf(time).multiply(ms).divide(BigDecimal.valueOf(count)).toString());
                }
            }
            Long execute = (Long) redisTemplate.execute(RATE_LIMIT, keys
                    , intervalPerTokens              //生成令牌的时间间隔 ms
                    , curTime                        //当前时间 默认
                    , rateLimit.bucketMaxTokens()    //桶初始化时的令牌数 默认为桶上限
                    , rateLimit.bucketMaxTokens()    //令牌桶的上限
                    , rateLimit.resetBucketInterval()//重置桶内令牌的时间间隔  默认为一天
            );
            log.info("LimitAspect---count:{}", execute);
            if (execute != null && execute.intValue() != 0) {
                log.info("LimitAspect---({})拿到令牌了,剩余令牌数量：{} 个", request.getRequestURI(), execute - 1);
                return joinPoint.proceed();
            } else {
                log.info("LimitAspect---({})没拿到令牌,限流啦~~", request.getRequestURI());
                throw new RuntimeException("接口" + request.getRequestURI() + "限流啦~");
            }
        } else {
            return joinPoint.proceed();
        }
    }

    public static String getIpAddr(HttpServletRequest request) {
        String ip = null;

        //X-Forwarded-For：Squid 服务代理
        String ipAddresses = request.getHeader("X-Forwarded-For");

        if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
            //Proxy-Client-IP：apache 服务代理
            ipAddresses = request.getHeader("Proxy-Client-IP");
        }

        if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
            //WL-Proxy-Client-IP：weblogic 服务代理
            ipAddresses = request.getHeader("WL-Proxy-Client-IP");
        }

        if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
            //HTTP_CLIENT_IP：有些代理服务器
            ipAddresses = request.getHeader("HTTP_CLIENT_IP");
        }

        if (ipAddresses == null || ipAddresses.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
            //X-Real-IP：nginx服务代理
            ipAddresses = request.getHeader("X-Real-IP");
        }

        //有些网络通过多层代理，那么获取到的ip就会有多个，一般都是通过逗号（,）分割开来，并且第一个ip为客户端的真实IP
        if (ipAddresses != null && ipAddresses.length() != 0) {
            ip = ipAddresses.split(",")[0];
        }

        //还是不能获取到，最后再通过request.getRemoteAddr();获取
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ipAddresses)) {
            ip = request.getRemoteAddr();
        }
        return ip;
    }
}
